#ifndef COMMON_H
#define COMMON_H

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <unordered_map>

using namespace nvonnxparser;
using namespace nvinfer1;




class Logger : public ILogger
{
    void  log(Severity severity, const char* msg) noexcept
    {
        if (severity <= Severity::kERROR)
        {
            std::cout << msg << std::endl;
        }
    }
};

enum trtModelType{
    YOLOV5 = 0,
    YOLOV8 = 1,
    YOLOV11 = 2,
    YOLOV5_SEG = 3,
    YOLOV8_SEG = 4,
    YOLOV11_SEG = 5,
    YOLOV8_POSE = 6,
    YOLOV11_POSE = 7,   
    YOLOV6_FACE = 8,
};

typedef struct trtModel{
    size_t size = static_cast<size_t>(-1);
    char *model{nullptr};
}trtModel;

// 模型配置结构体
struct ModelConfig {
    std::string model_path;
    std::string classes_path;
    std::string image_path;
    float conf_threshold;
    float nms_threshold;

    // 构造函数，提供默认值
    ModelConfig(
        const std::string& model = "",
        const std::string& classes = "",
        const std::string& image = "",
        float conf = 0.25f,
        float nms = 0.45f
    ) : model_path(model),
        classes_path(classes),
        image_path(image),
        conf_threshold(conf),
        nms_threshold(nms) {}
};

class ModelConfigManager {
public:
    static ModelConfigManager& getInstance() {
        static ModelConfigManager instance;
        return instance;
    }

    // 使用简单的文本格式加载配置
    void loadConfig(const std::string& config_path) {
        std::ifstream f(config_path);
        if (!f.is_open()) {
            throw std::runtime_error("Cannot open config file: " + config_path);
        }

        std::string line;
        std::string current_model;
        
        while (std::getline(f, line)) {
            if (line.empty() || line[0] == '#') continue;
            
            std::istringstream iss(line);
            std::string key, value;
            
            if (std::getline(iss, key, '=') && std::getline(iss, value)) {
                // 去除首尾空格
                key.erase(0, key.find_first_not_of(" \t"));
                key.erase(key.find_last_not_of(" \t") + 1);
                value.erase(0, value.find_first_not_of(" \t"));
                value.erase(value.find_last_not_of(" \t") + 1);
                
                if (key == "model_name") {
                    current_model = value;
                    configs[current_model] = ModelConfig();
                }
                else if (!current_model.empty()) {
                    if (key == "model_path") {
                        configs[current_model].model_path = value;
                    }
                    else if (key == "classes_path") {
                        configs[current_model].classes_path = value;
                    }else if (key == "image_path") {
                        configs[current_model].image_path = value;
                    }
                    else if (key == "conf_threshold") {
                        configs[current_model].conf_threshold = std::stof(value);
                    }
                    else if (key == "nms_threshold") {
                        configs[current_model].nms_threshold = std::stof(value);
                    }
                }
            }
        }
    }

    const ModelConfig& getConfig(const std::string& model_name) const {
        auto it = configs.find(model_name);
        if (it == configs.end()) {
            throw std::runtime_error("Model config not found: " + model_name);
        }
        return it->second;
    }

private:
    ModelConfigManager() {
        // 默认配置
        configs["yolov5"] = {
            "/path/to/yolov5.engine",
            "/path/to/classes.txt",
            "/path/to/image.jpg",
            0.25f,
            0.45f
        };
    }

    ModelConfigManager(const ModelConfigManager&) = delete;
    ModelConfigManager& operator=(const ModelConfigManager&) = delete;

    std::unordered_map<std::string, ModelConfig> configs;
};

#endif
